import re
import os
import sys
import json
from glob import glob

from readability.assert_rd import manage_assert, get_assert_summary

extern_func_dict = {}
func_dict = {}
extern_struct_and_type_dict = {}
struct_and_type_dict = {}
var_dict = {}

extracted_struct_and_type_str_dict = {}

extern_block_pattern = (r'extern\s\"C\"\s\{([^{}]*?)\}', 1)
extern_struct_and_type_pattern = (r'pub type (\w+);', 1)
extern_func_pattern = (r'fn (\w+)\(([^;]*?)\)([^;]*?);', 1)
func_pattern = (
    r'\n(\#\[no_mangle\]\n)?(pub\s)?(unsafe\s)?(extern\s\"C\"\s)?fn\s(\w+)\(([^{}]*?)\)([^{]*?)[\{(.*?)\n\}|\{\}]', 5)
struct_pattern = (r'\#\[derive\(Copy,\sClone\)\]\n\#\[repr\(C\)\]\npub\sstruct\s(\w+)\s\{(.*?)\n\}', 1)
type_pattern = (r'\npub\stype\s(\w+)\s\=\s([^;]*?);', 1)
var_pattern = (r'\n(#[no_mangle]\n)?(pub\s)?static\smut\s(\w+)\:\s(.*?)\s\=\s(unsafe\s)?([^;]*?);', 3)


def find_all_match(filename, content, pattern, pattern_dict, in_extern_block=False):
    pattern_dict.setdefault(filename, [])
    if not in_extern_block:
        extern_block_iters = re.finditer(extern_block_pattern[0], content, re.DOTALL)
        matches = re.finditer(pattern[0], content, re.DOTALL)
        for match_item in matches:
            pattern_dict[filename].append((match_item.group(pattern[1]), match_item.span()[0], match_item.span()[1]))
    else:
        extern_block_iters = re.finditer(extern_block_pattern[0], content, re.DOTALL)
        for extern_block_item in extern_block_iters:
            block_s, block_e = extern_block_item.span()
            block_content = content[block_s:block_e]
            matches = re.finditer(pattern[0], block_content, re.DOTALL)
            for match_item in matches:
                s, e = match_item.span()
                pattern_dict[filename].append((match_item.group(pattern[1]), s + block_s, e + block_s))


def find_dup_extern_func_by_filename(filename):
    result_index_set = set()
    extern_func_list_by_filename = extern_func_dict[filename]
    for i, (func_name, s, e) in enumerate(extern_func_list_by_filename):
        has_found = False
        # print(func_dict)
        for other_filename, other_func_list in func_dict.items():
            if filename == other_filename:
                continue
            for other_i, (other_func_name, other_s, other_e) in enumerate(other_func_list):
            # for other_i, item in enumerate(other_func_list):
            #     other_func_name, other_s, other_e = item
                if func_name == other_func_name:
                    result_index_set.add((i, other_filename, other_i))
                    has_found = True
                    break
            if has_found:
                break
    return result_index_set


def find_dup_local_func_by_filename(filename):
    result_index_set = set()
    func_list_by_filename = func_dict[filename]
    for i, (func_name, s, e) in enumerate(func_list_by_filename):
        has_found = False
        for other_filename, other_func_list in func_dict.items():
            if filename == other_filename:
                continue
            for other_i, (other_func_name, other_s, other_e) in enumerate(other_func_list):
                # if filename.find('tree') != -1 and other_filename.find('tree') != -1:
                #     print(filename, other_filename, func_name, other_func_name, func_name == other_func_name)
                #     input()
                if func_name == other_func_name:
                    result_index_set.add(i)
                    has_found = True
                    break
            if has_found:
                break
    # if filename.find('test_avl_tree') != -1 or filename.find('test_rb_tree') != -1:
    #   print(filename, [func_dict[filename][i] for i in result_index_set])
    return result_index_set


def find_dup_local_var_by_filename(filename):
    result_index_set = set()
    var_list_by_filename = var_dict[filename]
    for i, (var_name, s, e) in enumerate(var_list_by_filename):
        has_found = False
        for other_filename, other_var_list in var_dict.items():
            if filename == other_filename:
                continue
            for other_i, (other_var_name, other_s, other_e) in enumerate(other_var_list):
                if var_name == other_var_name:
                    result_index_set.add(i)
                    has_found = True
                    break
            if has_found:
                break
    return result_index_set


def record_dup(filename, content):
    find_all_match(filename, content, extern_struct_and_type_pattern, extern_struct_and_type_dict, in_extern_block=True)
    find_all_match(filename, content, extern_func_pattern, extern_func_dict, in_extern_block=True)
    find_all_match(filename, content, func_pattern, func_dict)
    find_all_match(filename, content, struct_pattern, struct_and_type_dict)
    find_all_match(filename, content, type_pattern, struct_and_type_dict)
    find_all_match(filename, content, var_pattern, var_dict)
    # if filename.find('test_set') != -1 or filename.find('test_hash_table') != -1:
    #     print(var_dict[filename])
    #     input()


def manage_dup(filename, content):
    extern_func_list = extern_func_dict[filename]
    var_list = var_dict[filename]
    local_func_list = func_dict[filename]
    dup_extern_in_block_str_list = []

    # 对extern函数：对于extern "C"块外部引入的情况，需要转换为use形式。
    use_dict = {}
    extern_dup_list = find_dup_extern_func_by_filename(filename)
    for i, other_filename, other_i in extern_dup_list:
        use_dict.setdefault(other_filename, set())
        use_dict[other_filename].add(extern_func_list[i][0])

    for i, other_filename, other_i in extern_dup_list:
        _, s, e = extern_func_list[i]
        dup_extern_in_block_str_list.append(content[s: e])

    # 对内含函数：添加后缀
    local_dup_list = find_dup_local_func_by_filename(filename)
    real_filename = os.path.basename(filename).split('.')[0]
    for i in local_dup_list:
        func_name, _, _ = local_func_list[i]
        content = re.sub(r'\b' + func_name + r'\b', func_name + '_' + real_filename, content)
        # content = content.replace(func_name, func_name + '_' + real_filename)

    # 对类型别名和结构体：先记录所有extern内外的名称，以待全部统一转换到一个定义文件中
    temp_struct_and_type_str_dict = {}
    use_dict.setdefault(os.path.join('src', 'struct_and_type.rs'), set())
    for extern_name, s, e in extern_struct_and_type_dict[filename]:
        dup_extern_in_block_str_list.append(content[s: e])
        temp_struct_and_type_str_dict[extern_name] = content[s: e]
        use_dict[os.path.join('src', 'struct_and_type.rs')].add(extern_name)
    for local_name, s, e in struct_and_type_dict[filename]:
        dup_extern_in_block_str_list.append(content[s: e])
        temp_struct_and_type_str_dict[local_name] = content[s: e]
        use_dict[os.path.join('src', 'struct_and_type.rs')].add(local_name)
    extracted_struct_and_type_str_dict.update(temp_struct_and_type_str_dict)

    # 对变量：对于重名变量，添加后缀
    var_dup_list = find_dup_local_var_by_filename(filename)
    for i in var_dup_list:
        var_name, _, _ = var_list[i]
        # print('dup var:', var_name)
        content = re.sub(r'\b' + var_name + r'\b', var_name + '_' + real_filename, content)
        # content = content.replace(var_name, var_name + '_' + real_filename)

    # 统一删除代码
    for extern_in_block_str in dup_extern_in_block_str_list:
        content = content.replace(extern_in_block_str, '')

    # 统一添加引入
    use_str = ''
    for other_filename, func_list in use_dict.items():
        use_path = '::'.join(other_filename.split('.')[0].split('/'))
        use_items = '{' + ', '.join(func_list) + '}'
        use_str += f'use crate::{use_path}::{use_items};\n'
    content = use_str + content
    return content


def write_global_struct_and_type_and_lib(rust_input_dir, rust_output_dir):
    with open(os.path.join(rust_output_dir, 'src', 'struct_and_type.rs'), 'w') as f:
        f.write('\n\n'.join(list(extracted_struct_and_type_str_dict.values())))
    with open(os.path.join(rust_input_dir, 'lib.rs'), 'r') as f:
        lib_content = f.read()
    insert_pos = lib_content.rfind('}')
    lib_content = (lib_content[:insert_pos] + '\n' +
                   'pub mod struct_and_type;' +
                   '\n' + lib_content[insert_pos:])
    with open(os.path.join(rust_output_dir, 'lib.rs'), 'w') as f:
        f.write(lib_content)


def cleaning(content):
    assert_fail_cleaning = "\s+fn\s__assert_fail\(\n\s+__assertion:\s\*const\slibc::c_char,\n\s+__file:\s\*const\slibc::c_char,\n\s+__line:\slibc::c_uint,\n\s+__function:\s\*const\slibc::c_char,\n\s+\)\s->\s!;"
    content = re.sub(assert_fail_cleaning, '', content)
    content = re.sub('\n+', '\n', content)
    return content


def main(rust_input_dir, rust_output_dir):
    input_file_list = glob(os.path.join(rust_input_dir, '**', '*.*'), recursive=True)
    file_dict = {}
    for filename in input_file_list:
        if os.path.isfile(filename):

            rel_dir_path = os.path.dirname(os.path.relpath(filename, rust_input_dir))
            os.makedirs(os.path.join(rust_output_dir, rel_dir_path), exist_ok=True)

            # if os.path.isfile(os.path.join(rust_input_dir, filename)):
            # print(filename)
            if filename.endswith('.rs') and filename.find('src') != -1:
                with open(filename, 'r') as f:
                    file_content = f.read()
                # print(os.path.relpath(filename, rust_input_dir))
                file_dict[os.path.relpath(filename, rust_input_dir)] = file_content
            else:
                with open(filename, 'rb') as f:
                    file_content = f.read()
                with open(os.path.join(rust_output_dir, os.path.relpath(filename, rust_input_dir)), 'wb') as f:
                    f.write(file_content)
    # print(file_dict.keys())

    for filename, file_content in file_dict.items():
        record_dup(filename, file_content)

    for filename, file_content in file_dict.items():
        file_content = manage_dup(filename, file_content)
        # print(os.path.join(rust_output_dir, filename), filename)
        file_content = manage_assert(file_content)
        file_content = cleaning(file_content)
        with open(os.path.join(rust_output_dir, filename), 'w') as f:
            f.write(file_content)

    write_global_struct_and_type_and_lib(rust_input_dir, rust_output_dir)

    process_summary_str = json.dumps(get_assert_summary(), indent=4, ensure_ascii=False)
    with open(os.path.join(rust_output_dir, 'assert_summary.json'), 'w', encoding='utf-8') as f:
        f.write(process_summary_str)


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])

